import torch.optim as optim
from SAMs.sam import SAM
from SAMs.psam import PSAM

Optimizers = {
    'SGD': optim.SGD,
    'Adam': optim.Adam,
    'SAM': SAM,
    'PSAM': PSAM
}

def get_optimizer(optimizer, net, lr, weight_decay, alpha=None):
    
    is_sam = (optimizer == 'SAM' or optimizer == 'PSAM')
    
    if is_sam:
        base_optimizer = optim.SGD
        # 用 SAM 包装基础优化器
        if optimizer == 'SAM':
            real_optimizer = SAM(net.parameters(), base_optimizer, lr=lr, weight_decay=weight_decay, momentum=0.9)
        elif optimizer == 'PSAM':
            real_optimizer = PSAM(net.parameters(), base_optimizer, lr=lr, weight_decay=weight_decay, alpha=alpha, momentum=0.9)
    else:
        if optimizer == 'SGDM':
            real_optimizer = optim.SGD(
                net.parameters(),
                lr=lr,
                momentum=0.9,
                weight_decay=weight_decay
            )
        else:
            real_optimizer = Optimizers[optimizer](
                net.parameters(),
                lr=lr,
                weight_decay=weight_decay
            )
        
    return real_optimizer, is_sam